from ultralytics import YOLO
import cv2
import os

# 检验detect是否正确


# Load a model
model = YOLO("yolov8x.pt")  # pretrained YOLOv8n model
model = YOLO(r"runs\detect\detect_x_1\train6\weights\best.pt")  # pretrained YOLOv8n model
dataPath = r'data\pose4\train'
image_path = os.path.join(dataPath, r'images')
label_path = os.path.join(r'data\pose_detect', r'labels1')

# Run batched inference on a list of images
results = model([r"data\pose4\train\images\000002.jpg"])  # return a list of Results objects
# results = model(source=image_path)  # return a list of Results objects

# Process results list
for result in results:
    boxes = result.boxes  # Boxes object for bounding box outputs
    masks = result.masks  # Masks object for segmentation masks outputs
    keypoints = result.keypoints  # Keypoints object for pose outputs
    probs = result.probs  # Probs object for classification outputs
    obb = result.obb  # Oriented boxes object for OBB outputs
    # h, w, c = img.shape
    img = cv2.imread(r"data\pose4\train\images\000002.jpg")

    # keypoints = result.keypoints
    h, w = boxes.orig_shape
    xywhn = boxes.xywhn[0]
    xn = xywhn[0]
    yn = xywhn[1]
    wn = xywhn[2]
    hn = xywhn[3]
    cx = xn * w
    cy = yn * h
    weight = wn * w
    height = hn * h
    xmin = cx - weight/2
    ymin = cy - height/2
    xmax = cx + weight/2
    ymax = cy + height/2
    print((xmin,ymin),(xmax,ymax))
    cv2.rectangle(img,(int(xmin),int(ymin)),(int(xmax),int(ymax)),(0,255,0),2)
    cv2.imshow('1',img)
    cv2.waitKey(0)
    print(boxes)
    # print(masks)
    # print(keypoints)
    # print(probs)
    # print(obb)
    # result.show()  # display to screen
    # result.save(filename="result.jpg")  # save to disk